# coding: utf-8 
"""
@Time    : 2024/8/26 14:25
@Author  : Y.H LEE
"""
"""
Visualize Genetic Algorithm to find the shortest path for travel sales problem.

Visit my tutorial website for more: https://mofanpy.com/tutorials/
"""
import matplotlib.pyplot as plt
import numpy as np

N_CITIES = 28  # DNA size
CROSS_RATE = 0.2    # 交叉概率
MUTATE_RATE = 0.02  # 变异概率
POP_SIZE = 500      # 个体数目
N_GENERATIONS = 1000     # 繁衍代数


class GA(object):
    def __init__(self, DNA_size, cross_rate, mutation_rate, pop_size):
        self.DNA_size = DNA_size
        self.cross_rate = cross_rate
        self.mutate_rate = mutation_rate
        self.pop_size = pop_size
        # 生成初始种群
        self.pop = np.vstack([np.random.permutation(DNA_size) for _ in range(pop_size)])
    #     vstack是垂直（按照行顺序）的把数组给堆叠起来
    #     permutation如果x是整数，则随机排列np.arange(x)。若果x是数组，对其复制之后再搅乱其元素。

    def translateDNA(self, DNA, city_position):  # get cities' coord in order
        line_x = np.empty_like(DNA, dtype=np.float64)
        line_y = np.empty_like(DNA, dtype=np.float64)
        # empty_like生成和已有数组相同大小，类型的数组
        # 生成的数组元素为随机数。
        for i, d in enumerate(DNA):
            # 对一个可遍历的数据对象(如列表、元组或字符串)，enumerate会将该数据对象组合为一个索引序列，同时列出数据和数据下标。
            # i-->index, d-->data 获取该DNA选择路径的城市坐标数组
            city_coord = city_position[d]
            # 找到相应的位置
            # 将相应的X,Y赋值给两列数组 -> 分别赋给x, y
            line_x[i, :] = city_coord[:, 0]
            line_y[i, :] = city_coord[:, 1]
        return line_x, line_y

    def get_fitness(self, line_x, line_y):
        # 生成一个数组，确认相邻点之间的距离
        total_distance = np.empty((line_x.shape[0],), dtype=np.float64)
        for i, (xs, ys) in enumerate(zip(line_x, line_y)):
            # i-->index, (xs, ys)-->(x,y)的位置（还要回到开始）
            xs, ys = np.append(xs, xs[0]), np.append(ys, ys[0])
            total_distance[i] = np.sum(np.sqrt(np.square(np.diff(xs)) + np.square(np.diff(ys))))
        #     diff-->差分，向量化计算，举个例子--> a = [1,3,4,0] b = [2, 9, 5, 7] np.square(np.diff(a)) + np.square(np.diff(b)
        #     -->[2, 1, -4]^2 + [7, -4, 2]^2
        #     放大差别
        fitness = np.exp(self.DNA_size * 2 / total_distance)
        return fitness, total_distance

    def select(self, fitness):
        """选择新一代"""
        idx = np.random.choice(np.arange(self.pop_size), size=self.pop_size, replace=True, p=fitness / fitness.sum())
        return self.pop[idx]

    def crossover(self, parent, pop):
        if np.random.rand() < self.cross_rate:
            i_ = np.random.randint(0, self.pop_size, size=1)  # select another individual from pop
            cross_points = np.random.randint(0, 2, self.DNA_size).astype(np.bool)  # choose crossover points
            # 0-->false 1-->true
            # print(cross_points)
            # print(~cross_points)
            keep_city = parent[~cross_points]  # find the city number
            # ~在此处是非的意思 ravel()为扁平化函数
            swap_city = pop[i_, np.isin(pop[i_].ravel(), keep_city, invert=True)]
            # np.isin(a,b) 用于判定a中的元素在b中是否出现过，如果出现过返回True,否则返回False,最终结果为一个形状和a一模一样的数组。
            # 但是当参数invert被设置为True时，情况恰好相反，如果a中元素在b中没有出现则返回True,如果出现了则返回False.
            parent[:] = np.concatenate((keep_city, swap_city))
        return parent

    def mutate(self, child):
        # 在 mutate 的时候, 也是找到两个不同的 DNA 点, 然后交换这两个点就好了.
        for point in range(self.DNA_size):
            if np.random.rand() < self.mutate_rate:
                swap_point = np.random.randint(0, self.DNA_size)
                swapA, swapB = child[point], child[swap_point]
                child[point], child[swap_point] = swapB, swapA
        return child

    def evolve(self, fitness):
        # 选择
        pop = self.select(fitness)
        pop_copy = pop.copy()   # 备个份
        for parent in pop:  # for every parent
            # 这里parent为遍历pop，一次为其中一行，而这里的pop是从原pop中按适应度概率有放回的选出了POP_SIZE行
            # 交叉
            child = self.crossover(parent, pop_copy)
            # 变异
            child = self.mutate(child)
            parent[:] = child
            # print(parent[:])
            # print(parent)
        #     两者是一致的
        self.pop = pop


class TravelSalesPerson(object):
    def __init__(self, n_cities):
        self.city_position = np.random.rand(n_cities, 2)
        plt.switch_backend('TkAgg')
        plt.ion()

    def plotting(self, lx, ly, total_d):
        lx, ly = np.append(lx, lx[0]), np.append(ly, ly[0])
        plt.cla()
        # 画城市所在的点
        plt.scatter(self.city_position[:, 0].T, self.city_position[:, 1].T, s=100, c='k')
        # 画路线
        plt.plot(lx.T, ly.T, 'r-')
        plt.text(-0.05, -0.05, "Total distance=%.2f" % total_d, fontdict={'size': 20, 'color': 'red'})
        plt.xlim((-0.1, 1.1))
        plt.ylim((-0.1, 1.1))
        plt.pause(0.01)


# 实例化对象
ga = GA(DNA_size=N_CITIES, cross_rate=CROSS_RATE, mutation_rate=MUTATE_RATE, pop_size=POP_SIZE)
# 实例化地图
env = TravelSalesPerson(N_CITIES)
for generation in range(N_GENERATIONS):
    lx, ly = ga.translateDNA(ga.pop, env.city_position)
    fitness, total_distance = ga.get_fitness(lx, ly)
    ga.evolve(fitness)
    best_idx = np.argmax(fitness)
    print('Gen:', generation, '| best fit: %.2f' % fitness[best_idx], )

    env.plotting(lx[best_idx], ly[best_idx], total_distance[best_idx])

plt.ioff()
plt.show()
